"""MCP Client Manager Module"""
import json
from typing import Dict, Any

from mcp.types import CallToolResult

from .connection import ServerConnection
from .config_model import MCPConfig
from loguru import logger
import asyncio

class MCPClientManager:
    """Main client manager that coordinates all components"""

    def __init__(self, config: MCPConfig | str):
        # 处理配置参数
        if isinstance(config, str):
            if config.endswith('.json'):
                self.config = MCPConfig.from_json_file(config)
            else:
                config_data = json.loads(config)
                self.config = MCPConfig(**config_data)
        else:
            # 直接使用提供的MCPConfig对象
            self.config = config

        self.connections: Dict[str, ServerConnection] = {}

        # Initialize connections dict but don't connect yet
        for server_type in self.config.mcpServers:
            self.connections[server_type] = None

    async def initialize(self):
        """Initialize connections to all MCP servers"""
        for server_type in self.config.mcpServers:
            self.connections[server_type] = await self.create_session(server_type)
            logger.info(f"Successfully connected to {server_type} server")

    async def create_session(self, server_type: str) -> ServerConnection:
        """Create a new MCP session"""
        if server_type not in self.config.mcpServers:
            logger.error(f"Invalid server type: {server_type}")
            raise ValueError(f"Invalid server type: {server_type}")

        server_config = self.config.mcpServers[server_type]
        connection = ServerConnection(
            name=server_type,
            config={
                "command": server_config.command,
                "args": server_config.args,
                "retries": 3,
                "retry_delay": 1.0
            }
        )

        await connection.connect()
        if not connection.is_connected():
            logger.error(f"Failed to connect to {server_type} server")
            raise ConnectionError(f"Failed to connect to {server_type} server")
        logger.info(f"Created new session for {server_type} server")
        return connection

    async def destroy_session(self, server_type: str) -> None:
        """Destroy an existing MCP session"""
        if server_type in self.connections:
            await self.connections[server_type].cleanup()
            del self.connections[server_type]
            if not self.connections:
                self.tool_manager = None
            logger.info(f"Destroyed session for {server_type} server")

    async def cleanup(self) -> None:
        """Cleanup all resources"""
        await asyncio.gather(*[
            connection.cleanup() for connection in self.connections.values()
        ])
        self.connections.clear()
        logger.info("Cleaned up all resources")

    def is_available(self) -> bool:
        """Check if MCPClient is available
        """
        if not self.connections:
            return False

        for connection in self.connections.values():
            if connection and connection.is_connected():
                return True

        return False

    async def get_tools(self, server_type: str) -> list:
        """Get tool list for specific server"""
        if server_type not in self.connections:
            logger.error(f"Invalid server type: {server_type}")
            raise ValueError(f"Invalid server type: {server_type}")

        try:
            tools = await self.connections[server_type].tool_manager.list_tools()
            return tools
        except Exception as e:
            logger.error(f"Failed to get tools for {server_type}: {str(e)}")
            # Try to reconnect and get tools again
            logger.info(f"Attempting to reconnect to {server_type} server")
            self.connections[server_type] = await self.create_session(server_type)
            return await self.connections[server_type].tool_manager.list_tools()

    async def get_all_tools(self) -> Dict[str, list]:
        """Get tool list for all servers"""
        all_tools = {}
        for server_type, conn in self.connections.items():
            try:
                all_tools[server_type] = await conn.tool_manager.list_tools()
            except Exception as e:
                logger.error(f"Failed to get tools for {server_type}: {str(e)}")
                # Try to reconnect and get tools again
                logger.info(f"Attempting to reconnect to {server_type} server")
                self.connections[server_type] = await self.create_session(server_type)
                all_tools[server_type] = await self.connections[server_type].tool_manager.list_tools()
        return all_tools

    async def get_all_tools_str(self) -> str:
        """Get all tools as a formatted string"""
        tools_dict = await self.get_all_tools()
        result = []
        for server, tools in tools_dict.items():
            tools_str = "\n".join([tool.format_for_llm() for tool in tools])
            result.append(tools_str)
        return "----------".join(result)

    async def execute_tool(self, server_type: str, tool_name: str, arguments: Dict[str, Any]) -> CallToolResult:
        """Execute tool on specific server"""
        if server_type not in self.connections:
            logger.error(f"Invalid server type: {server_type}")
            raise ValueError(f"Invalid server type: {server_type}")

        try:
            result = await self.connections[server_type].tool_manager.execute_tool(tool_name, arguments)
            return result
        except Exception as e:
            logger.error(f"Failed to execute tool {tool_name} on {server_type}: {str(e)}")
            # Try to reconnect and execute tool again
            logger.info(f"Attempting to reconnect to {server_type} server")
            self.connections[server_type] = await self.create_session(server_type)
            result = await self.connections[server_type].tool_manager.execute_tool(tool_name, arguments)
            logger.info(f"Successfully executed tool {tool_name} on {server_type} after reconnection")
            return result
